'''
@Time    : 2022/3/18 18:45
@Author  : Fu Junyu
@Site    : www.fujunyu.cn
@File    : test.py
@Software: PyCharm
'''
import dgl
import numpy as np
import torch
from dgl.nn.pytorch import GraphConv
from torch import nn
import torch.nn.functional as F
import networkx as nx
import matplotlib.pyplot as plt
import dgl.function as fn

'''
传统同构图（Homogeneous Graph）数据中只存在一种节点和边，因此在构建图神经网络时所有节点共享同样的模型参数并且拥有同样维度的特征空间。
而异构图（Heterogeneous Graph）中可以存在不只一种节点和边，因此允许不同类型的节点拥有不同维度的特征或属性。
'''

trainUser = []
trainItem = []
trainRating = []

train_file = '../data/train.txt'
with open(train_file) as f:
    for l in f.readlines():
        if len(l) > 0:
            l = l.strip('\n').split(' ')
            items = int(l[1])
            rating = float(str(l[2]))
            uid = int(l[0])

            trainRating.append(rating)
            trainUser.append(uid)
            trainItem.append(items)

trainUser = torch.LongTensor(np.array(trainUser))
trainItem = torch.LongTensor(np.array(trainItem))
trainRating = torch.FloatTensor(np.array(trainRating))


graph_data = {
    ('user', 'rating', 'movies'): (trainUser, trainItem),
}

g = dgl.heterograph(graph_data)

# 查看节点和边
# print(g.nodes('user'))
# print(g.nodes('movies'))
# print(g.edges(etype='rating'))

# 设置节点和边的属性
g.nodes['user'].data['x'] = torch.randn(g.num_nodes('user'), 4)
g.nodes['movies'].data['y'] = torch.randn(g.num_nodes('movies'), 4)
g.edges['rating'].data['w'] = trainRating.reshape(trainRating.shape, -1)

g.nodes['user'].data['y'] = torch.randn(g.num_nodes('user'), 4)
g.nodes['movies'].data['x'] = torch.randn(g.num_nodes('movies'), 4)
g.nodes['user'].data['tf'] = torch.randn(g.num_nodes('user'), 4)
g.nodes['movies'].data['tf'] = torch.randn(g.num_nodes('movies'), 4)


etypes = g.canonical_etypes  # 返回源节点类型、边缘类型和目标节点类型的字符串三元组(str、str、str)

# for src, edg, dst in etypes:
#     print(src, edg, dst)
# print(etypes)

class RgcnLayer(nn.Module):
    def __init__(self, in_size, out_size, etypes):
        super(RgcnLayer, self).__init__()
        # 定义关系的参数矩阵W_r
        self.weight = nn.ModuleDict({
            name: nn.Linear(in_size, out_size) for name in etypes
        })

    def forward(self, G, fea_dict):
        func = {}

        for srctype, etype, dsttype in G.canonical_etypes:
            # 计算W_r * h
            print(srctype)
            print('fea_dict[srctype]:', fea_dict[srctype])
            Wh = self.weight[etype](fea_dict[srctype])

            G.nodes[srctype].data['Wh_%s' % etype] = Wh

            # print('Wh_rating:', G.nodes[srctype].data['Wh_rating'])

            func[etype] = (fn.copy_u('Wh_%s' % etype, 'm'), fn.mean('m', 'tf'))

            # print(func)

        G.multi_update_all(func, 'sum')
        # print(G.nodes['movies'].data['tf'])

        return {ntype: G.nodes[ntype].data['tf'] for ntype in G.ntypes}


# nid = g.subgraph([1,2])
# print(nid)

# test = g.nodes().mailbox['m']
# print(test)


u, v = torch.tensor([0, 0, 1, 1, 2, 3, 3]), torch.tensor([0, 1, 1, 1, 2, 3, 4])
graph_data = {
    ('user', 'rating', 'movies'): (u, v),
    ('movies', 'rec', 'user'): (v, u),
}
g = dgl.heterograph(graph_data)
# print(g.nodes('user'))
# print(g.edges())

g.nodes['user'].data['x'] = torch.ones(g.num_nodes('user'), 4)
g.nodes['movies'].data['y'] = torch.randn(g.num_nodes('movies'), 4)


g.nodes['user'].data['y'] = torch.randn(g.num_nodes('user'), 4)
g.nodes['movies'].data['x'] = torch.zeros(g.num_nodes('movies'), 4)
g.nodes['user'].data['tf'] = torch.ones(g.num_nodes('user'), 4)
g.nodes['movies'].data['tf'] = torch.ones(g.num_nodes('movies'), 4)


def update_all_example(graph):
    # 在graph.ndata['ft']中存储结果
    graph.update_all(fn.u_add_v('tf', 'tf', 'h'),
                     fn.sum('h', 'tf'))
    # 在update_all外调用更新函数
    # final_ft = graph.ndata['ft'] * 2
    # graph.ndata['tf'] = final_ft
    # return final_ft
    print(graph.ndata['tf'])
# print(g.ndata['x'])
# print(g.ndata['y'])
# print(g.ndata['tf'])
# update_all_example(g)

# layer = RgcnLayer(4, 3, ['rating', 'rec'])
#
#
# print(layer(g, g.ndata['x']))


graph_data = {
    ('user', 'rating', 'user'): (torch.tensor([0, 1]), torch.tensor([1, 1])),
    ('movies', 'rec', 'user'):(torch.tensor([0]), torch.tensor([1]))
}

g = dgl.heterograph(graph_data)

g.nodes['user'].data['h'] = torch.tensor([1., 2.])
g.nodes['movies'].data['h'] = torch.tensor([1.])
# print(g.nodes['user'])

# g.nodes['user'].data['h'] = torch.ones(g.num_nodes('user'), 2)
# g.nodes['movies'].data['h'] =torch.ones(g.num_nodes('movies'), 2)

print(g.nodes['user'].data['h'])



# g.multi_update_all({
#     'rating': (fn.u_mul_v('h','h',  'm'), fn.sum('m', 'h')),
#     'rec': (fn.u_mul_v('h','h', 'm'), fn.sum('m', 'h')),
# }, 'sum')

# print(g.nodes['user'].data['h'])
# print(g.nodes['movies'].data['h'])

print(g.ndata)

